// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/init/init.h>
#include <turbo/flags/flag.h>
#include <gtest/gtest.h>
#include <unordered_set>

#include <pollux/testing/exec/aggregation_fuzzer_options.h>
#include <pollux/testing/exec/aggregation_fuzzer_runner.h>
#include <pollux/testing/exec/presto_query_runner.h>
#include <pollux/testing/exec/transform_result_verifier.h>
#include <pollux/functions/prestosql/aggregates/register_aggregate_functions.h>
#include <pollux/functions/prestosql/fuzzer/ApproxDistinctInputGenerator.h>
#include <pollux/functions/prestosql/fuzzer/ApproxDistinctResultVerifier.h>
#include <pollux/functions/prestosql/fuzzer/ApproxPercentileInputGenerator.h>
#include <pollux/functions/prestosql/fuzzer/ApproxPercentileResultVerifier.h>
#include <pollux/functions/prestosql/fuzzer/ArbitraryResultVerifier.h>
#include <pollux/functions/prestosql/fuzzer/AverageResultVerifier.h>
#include <pollux/functions/prestosql/fuzzer/ClassificationAggregationInputGenerator.h>
#include <pollux/functions/prestosql/fuzzer/MapUnionSumInputGenerator.h>
#include <pollux/functions/prestosql/fuzzer/MinMaxByResultVerifier.h>
#include <pollux/functions/prestosql/fuzzer/MinMaxInputGenerator.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/functions/prestosql/window/window_functions_registration.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_FLAG(int64_t,
    seed,
    0,
    "Initial seed for random number generator used to reproduce previous "
    "results (0 means start with random seed).");

TURBO_FLAG(std::string,
    only,
    "",
    "If specified, Fuzzer will only choose functions from "
    "this comma separated list of function names "
    "(e.g: --only \"min\" or --only \"sum,avg\").");

TURBO_FLAG(std::string,
    presto_url,
    "",
    "Presto coordinator URI along with port. If set, we use Presto "
    "source of truth. Otherwise, use DuckDB. Example: "
    "--presto_url=http://127.0.0.1:8080");

TURBO_FLAG(uint32_t,
    req_timeout_ms,
    1000,
    "Timeout in milliseconds for HTTP requests made to reference DB, "
    "such as Presto. Example: --req_timeout_ms=2000");

// Any change made in the file should be reflected in
// the FB-internal aggregation fuzzer test too.

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
namespace {

std::unordered_map<std::string, std::shared_ptr<InputGenerator>>
getCustomInputGenerators() {
  return {
      {"min", std::make_shared<MinMaxInputGenerator>("min")},
      {"min_by", std::make_shared<MinMaxInputGenerator>("min_by")},
      {"max", std::make_shared<MinMaxInputGenerator>("max")},
      {"max_by", std::make_shared<MinMaxInputGenerator>("max_by")},
      {"approx_distinct", std::make_shared<ApproxDistinctInputGenerator>()},
      {"approx_set", std::make_shared<ApproxDistinctInputGenerator>()},
      {"approx_percentile", std::make_shared<ApproxPercentileInputGenerator>()},
      {"map_union_sum", std::make_shared<MapUnionSumInputGenerator>()},
      {"classification_fall_out",
       std::make_shared<ClassificationAggregationInputGenerator>()},
      {"classification_precision",
       std::make_shared<ClassificationAggregationInputGenerator>()},
      {"classification_recall",
       std::make_shared<ClassificationAggregationInputGenerator>()},
      {"classification_miss_rate",
       std::make_shared<ClassificationAggregationInputGenerator>()},
      {"classification_thresholds",
       std::make_shared<ClassificationAggregationInputGenerator>()}};
}

} // namespace
} // namespace kumo::pollux::exec::test

int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);

  // Calls common init functions in the necessary order, initializing
  // singletons, installing proper signal handlers for better debugging
  // experience
  melon::Init init(&argc, &argv);

  // Register only presto supported signatures if we are verifying against
  // Presto.
  if (FLAGS_presto_url.empty()) {
    kumo::pollux::aggregate::prestosql::registerAllAggregateFunctions(
        "", false);
  } else {
    kumo::pollux::aggregate::prestosql::registerAllAggregateFunctions(
        "", false, true);
  }

  kumo::pollux::functions::prestosql::registerAllScalarFunctions();
  kumo::pollux::window::prestosql::registerAllWindowFunctions();
  kumo::pollux::functions::prestosql::registerInternalFunctions();
  kumo::pollux::aggregate::prestosql::registerInternalAggregateFunctions();
  kumo::pollux::memory::MemoryManager::initialize({});

  size_t initialSeed = FLAGS_seed == 0 ? std::time(nullptr) : FLAGS_seed;

  // List of functions that have known bugs that cause crashes or failures.
  static const std::unordered_set<std::string> skipFunctions = {
      // Skip internal functions used only for result verifications.
      "$internal$count_distinct",
      "$internal$array_agg",
      // https://github.com/facebookincubator/pollux/issues/3493
      "stddev_pop",
      // Lambda functions are not supported yet.
      "reduce_agg",
      "max_data_size_for_stats",
      "any_value",
  };

  static const std::unordered_set<std::string> functionsRequireSortedInput = {
      "tdigest_agg",
  };

  using kumo::pollux::exec::test::ApproxDistinctResultVerifier;
  using kumo::pollux::exec::test::ApproxPercentileResultVerifier;
  using kumo::pollux::exec::test::ArbitraryResultVerifier;
  using kumo::pollux::exec::test::AverageResultVerifier;
  using kumo::pollux::exec::test::MinMaxByResultVerifier;
  using kumo::pollux::exec::test::setupReferenceQueryRunner;
  using kumo::pollux::exec::test::TransformResultVerifier;

  auto makeArrayVerifier = []() {
    return TransformResultVerifier::create("\"$internal$canonicalize\"({})");
  };

  auto makeMapVerifier = []() {
    return TransformResultVerifier::create(
        "\"$internal$canonicalize\"(map_keys({}))");
  };

  // Functions whose results verification should be skipped. These can be
  // order-dependent functions whose results depend on the order of input rows,
  // or functions that return complex-typed results containing floating-point
  // fields. For some functions, the result can be transformed to a value that
  // can be verified. If such transformation exists, it can be specified to be
  // used for results verification. If no transformation is specified, results
  // are not verified.
  static const std::unordered_map<
      std::string,
      std::shared_ptr<kumo::pollux::exec::test::ResultVerifier>>
      customVerificationFunctions = {
          // Order-dependent functions.
          {"approx_distinct", std::make_shared<ApproxDistinctResultVerifier>()},
          {"approx_set", std::make_shared<ApproxDistinctResultVerifier>(true)},
          {"approx_percentile",
           std::make_shared<ApproxPercentileResultVerifier>()},
          {"arbitrary", std::make_shared<ArbitraryResultVerifier>()},
          {"any_value", nullptr},
          {"array_agg", makeArrayVerifier()},
          {"set_agg", makeArrayVerifier()},
          {"set_union", makeArrayVerifier()},
          {"map_agg", makeMapVerifier()},
          {"map_union", makeMapVerifier()},
          {"map_union_sum", makeMapVerifier()},
          {"max_by", std::make_shared<MinMaxByResultVerifier>(false)},
          {"min_by", std::make_shared<MinMaxByResultVerifier>(true)},
          {"avg", std::make_shared<AverageResultVerifier>()},
          {"multimap_agg",
           TransformResultVerifier::create(
               "transform_values({}, (k, v) -> \"$internal$canonicalize\"(v))")},
          // Semantically inconsistent functions
          {"skewness", nullptr},
          {"kurtosis", nullptr},
          {"entropy", nullptr},
          // https://github.com/facebookincubator/pollux/issues/6330
          {"max_data_size_for_stats", nullptr},
          {"sum_data_size_for_stats", nullptr},
      };

  using Runner = kumo::pollux::exec::test::AggregationFuzzerRunner;
  using Options = kumo::pollux::exec::test::AggregationFuzzerOptions;

  Options options;
  options.onlyFunctions = FLAGS_only;
  options.skipFunctions = skipFunctions;
  options.functionsRequireSortedInput = functionsRequireSortedInput;
  options.customVerificationFunctions = customVerificationFunctions;
  options.customInputGenerators =
      kumo::pollux::exec::test::getCustomInputGenerators();
  options.timestampPrecision =
  options.timestampPrecision =
      kumo::pollux::VectorFuzzer::Options::TimestampPrecision::kMilliSeconds;
  std::shared_ptr<kumo::pollux::memory::MemoryPool> rootPool{
      kumo::pollux::memory::memoryManager()->addRootPool()};

  return Runner::run(
      initialSeed,
      setupReferenceQueryRunner(
          rootPool.get(),
          FLAGS_presto_url,
          "aggregation_fuzzer",
          FLAGS_req_timeout_ms),
      options);
}
